from PyQt5.QtCore import QAbstractTableModel, Qt, pyqtSignal, QModelIndex

from pandas import DataFrame

from utils.define import UNSHOW_INDEX

class PandasModel(QAbstractTableModel):
    dataChanged = pyqtSignal()

    def __init__(self, data:DataFrame):
        QAbstractTableModel.__init__(self)
        self._data = data
        # 过滤出未被标记为删除的数据(filtered_data 中 index字段的行号会缺失 被过滤数据的行号)
        self.base_filtered = None
        self.filtered_data = None
        self.currentYM = None # 结构(2023, 1)
        self.refreshFilterData()

        self.original_columns = list(data.columns)
        self.display_columns = [col for col in self.original_columns if col != 'deleted']

    def refreshFilterData(self):
        self.base_filtered = self._data[self._data['deleted'] != True].drop(columns=UNSHOW_INDEX)

        # 如果设置了当前月份，则进一步过滤
        if self.currentYM is not None:
            # 筛选指定年月的数据
            year, month = self.currentYM
            year_month_prefix = f"{year}.{month}."
            month_filtered = self.base_filtered[
                self.base_filtered['日期'].astype(str).str.startswith(year_month_prefix)
            ]
            self.filtered_data = month_filtered
        else:
            self.filtered_data = self.base_filtered

    def setCurrentYM(self, year_month):
        # 设置当前显示的年月，None表示显示所有数据
        self.currentYM = year_month
        self.refreshFilterData()
        self.layoutChanged.emit()  # 通知视图重新布局

    def actualRow(self, row):
        return self.filtered_data.index[row]

    def actualRowAndCol(self, index: QModelIndex):
        row = self.filtered_data.index[index.row()]
        col = self.original_columns.index(self.display_columns[index.column()])
        return row, col

    def sumAll(self):
        result = 0
        for price in self.base_filtered["金额"]:
            result += float(price)
        return "{:.2f}".format(result)

    def sumMonth(self, year_month):
        result = 0
        year, month = year_month[0], year_month[1]
        for i in self.filtered_data.index:
            date = str(self.filtered_data.loc[i, "日期"]).split(".")
            dateY, dateM = date[0], date[1]
            if dateY == year and dateM == month:
                result += float(self.filtered_data.loc[i, "金额"])
        return "{:.2f}".format(result)

    def rowCount(self, parent=None):
        return self.filtered_data.shape[0]

    def columnCount(self, parent=None):
        return self.filtered_data.shape[1]

    def data(self, index, role=Qt.DisplayRole):
        if index.isValid():
            if role == Qt.DisplayRole or role == Qt.EditRole:
                return str(self._data.iloc[self.actualRowAndCol(index)])
        return None

    def setData(self, index, value, role=Qt.EditRole):
        if index.isValid() and role == Qt.EditRole and value != "nan":
            self._data.iloc[self.actualRowAndCol(index)] = value
            # 重新计算过滤后的数据
            self.refreshFilterData()
            # 发送信号的index是过滤数据的 以后要用到需要修改并需要重新计算
            self.dataChanged.emit()
            return True
        return False

    def headerData(self, col, orientation, role=Qt.DisplayRole):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return self._data.columns[col]
        return None

    def removeRow(self, row):
        self._data.loc[self.actualRow(row), 'deleted'] = True
        self.dataChanged.emit()
        self.refreshFilterData()
        return True

    # TODO: 修改数据限制
    # def flags(self, index):
    #     return QAbstractTableModel.flags(self, index) | Qt.ItemIsEditable